package main

import (
	"flag"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	var directory string
	var verbose bool
	flag.StringVar(&directory, "dir", ".", "directory to search for pb.go files")
	flag.BoolVar(&verbose, "verbose", false, "verbosity")
	flag.Parse()
	filepathErr := filepath.Walk(directory, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}
		// ignore all files that have not been generated by pb
		if strings.HasSuffix(path, ".pb.go") == false {
			return nil
		}
		originalSize := info.Size()
		fset := token.NewFileSet()
		node, err := parser.ParseFile(fset, path, nil, parser.ParseComments)
		if err != nil {
			return fmt.Errorf("could not parse go file: %v", err)
		}
		var v visitor
		// modify tags
		ast.Walk(v, node)
		// output file
		f, err := os.Create(path)
		if err != nil {
			return fmt.Errorf("could not open output file: %v", err)
		}
		defer f.Close()
		err = format.Node(f, fset, node)
		if err != nil {
			return fmt.Errorf("could not write output file: %v", err)
		}
		if verbose {
			newInfo, err := os.Stat(path)
			if err != nil {
				fmt.Println(err)
			}
			if originalSize != newInfo.Size() {
				fmt.Print("modified:")
			}
			fmt.Println(path)
		}
		return nil
	})
	if filepathErr != nil {
		fmt.Printf("could not correctly handle directory: %v", filepathErr)
	}
}

// visitor to handle every ast node
type visitor struct{}

// Visit find fields starting with XXX_ and add bson tag
func (v visitor) Visit(n ast.Node) ast.Visitor {
	if n == nil {
		return nil
	}
	switch n.(type) {
	case *ast.Field:
		ast.Inspect(n, func(subn ast.Node) bool {
			switch subd := subn.(type) {
			case *ast.Field:
				if f, ok := n.(*ast.Field); ok {
					return len(f.Names) > 0
				}
			case *ast.BasicLit:
				existingTagValue := strings.Trim(subd.Value, "`")
				bsonTags := []string{"bson:\"-\""}
				if existingTagValue == "" {
					subd.Value = fmt.Sprintf("`%s`", strings.Join(bsonTags, " "))
					return true
				}
				// add or replace additional tags
				addBsonTags := []string{}
				existingTags := strings.Split(existingTagValue, " ")
				for _, additionalTag := range bsonTags {
					additionalTagParts := strings.Split(additionalTag, ":")
					foundIndex := -1
					for index, existingTag := range existingTags {
						if strings.HasPrefix(existingTag, "json:") {
							tmpTag := strings.Replace(existingTag, "json", "bson", 1)
							overTag := strings.Replace(tmpTag, ",omitempty", "", 1)
							addBsonTags = append(addBsonTags, overTag)
						}
						existingTagParts := strings.Split(existingTag, ":")
						if fmt.Sprintf("%v:", additionalTagParts[0]) == fmt.Sprintf("%v:", existingTagParts[0]) {
							foundIndex = index
							break
						}
					}
					if foundIndex > -1 {
						existingTags = append(existingTags[:foundIndex], existingTags[foundIndex+1:]...)
					}
				}
				mergedTags := append(existingTags, addBsonTags...)
				subd.Value = fmt.Sprintf("`%s`", strings.Join(mergedTags, " "))
			}
			return true
		})
	}
	return v
}
